__device__
float2 complex_new()
{
	return make_float2(0, 0);
}

__device__
float2 copy(float2 arg)
{
	return make_float2(arg.x, arg.y);
}

__device__
float2 conjugate(float2 arg)
{
	return make_float2(arg.x, -arg.y);
}

__device__
float2 complex_exp(float arg)
{
	return make_float2(cosf(arg), sinf(arg));
}

__device__
float2 complex_add(float2 a, float2 b)
{
	return make_float2(a.x + b.x, a.y + b.y);
}

__device__
float2 complex_sub(float2 a, float2 b){
	return make_float2(a.x - b.x, a.y - b.y);
}

__device__
float2 complex_multy(float2 a, float2 b){
	return make_float2(a.x*b.x - a.y*b.y, a.y*b.x + a.x*b.y);
}

__device__ __constant__ float _PI = 3.141592653;

extern "C"
__global__ void fft(int size, float2* a, float2* wlist, int iteration){
	int idx = blockIdx.x * gridDim.x + threadIdx.x;
	int idy = blockIdx.y * gridDim.y + threadIdx.y;
	if (idx < size && idy < size){
		int step = 1 << iteration;
		int step2 = step*2;
		bool flag = (idy/step) % 2 + (idx/step) % 2 > 0;
		if (!flag){
			float2 w1 = wlist[(idy % step2)*(size/step2)];
			float2 w2 = wlist[(idx % step2)*(size/step2)];
			
			int a00i = idy*size + idx;
			int a10i = (idy+step)*size + idx;
			int a01i = idy*size + (idx+step);
			int a11i = (idy+step)*size + (idx+step);
			float2 a00 = a[a00i];
			float2 a10 = a[a10i];
			float2 a01 = a[a01i];
			float2 a11 = a[a11i];
			
			float2 newa00 = complex_new();
			float2 newa01 = complex_new();
			float2 newa10 = complex_new();
			float2 newa11 = complex_new();
			
			newa00 = complex_add(newa00, a00);
			newa10 = complex_add(newa10, a00);
			newa01 = complex_add(newa01, a00);
			newa11 = complex_add(newa11, a00);
			
			float2 t = complex_multy(w1, a10);
			newa00 = complex_add(newa00, t);
			newa10 = complex_sub(newa10, t);
			newa01 = complex_add(newa01, t);
			newa11 = complex_sub(newa11, t);
			
			t = complex_multy(w2, a01);
			newa00 = complex_add(newa00, t);
			newa10 = complex_add(newa10, t);
			newa01 = complex_sub(newa01, t);
			newa11 = complex_sub(newa11, t);
			
			t = complex_multy(w2, a11);
			t = complex_multy(w1, t);
			newa00 = complex_add(newa00, t);
			newa10 = complex_sub(newa10, t);
			newa01 = complex_sub(newa01, t);
			newa11 = complex_add(newa11, t);
			
			a[a00i] = newa00;
			a[a10i] = newa10;
			a[a01i] = newa01;
			a[a11i] = newa11;
		}
	}
	__syncthreads();
}

extern "C"
__global__ void fftRows(int size, float2* src, float2* wlist){
	int blockSize = blockDim.x * blockDim.y;
	int block = blockIdx.y * gridDim.x + blockIdx.x;
	int idx = (blockSize * block) + threadIdx.y * blockDim.x + threadIdx.x;
	int row = idx;
	if (idx < size){
		for (int q = 0; q < size; q++) {
			int step = 1 << q;
			int step2 = step*2;
			for (int i = 0; i < size - step; i++) {
				bool flag = (i/step) % 2;
				if (!flag){
					float2 w = wlist[(i % step2)*(size/step2)];
					int ai = row*size + i;
					int bi = row*size + i + step;
					float2 a = src[ai];
					float2 b = src[bi];
					
					float2 t = complex_multy(b, w);
					float2 anew = complex_add(a, t);
					float2 bnew = complex_sub(a, t);
					
					src[ai] = anew;
					src[bi] = bnew;
				}
			}
		}	
	}
	__syncthreads();
}

extern "C"
__global__ void fftCols(int size, float2* src, float2* wlist){
	int blockSize = blockDim.x * blockDim.y;
	int block = blockIdx.y * gridDim.x + blockIdx.x;
	int idx = (blockSize * block) + threadIdx.y * blockDim.x + threadIdx.x;
	int col = idx;
	if (idx < size){
		for (int q = 0; q < size; q++) {
			int step = 1 << q;
			int step2 = step*2;
			for (int i = 0; i < size - step; i++) {
				bool flag = (i/step) % 2;
				if (!flag){
					float2 w = wlist[(i % step2)*(size/step2)];
					int ai = i*size + col;
					int bi = (i + step)*size + col;
					float2 a = src[ai];
					float2 b = src[bi];
					
					float2 t = complex_multy(b, w);
					float2 anew = complex_add(a, t);
					float2 bnew = complex_sub(a, t);
					
					src[ai] = anew;
					src[bi] = bnew;
				}
			}
		}
	}
	__syncthreads();
}

extern "C"
__global__ void wlist(int size, float2* w){
	int blockSize = blockDim.x * blockDim.y;
	int block = blockIdx.y * gridDim.x + blockIdx.x;
	int idx = (blockSize * block) + threadIdx.y * blockDim.x + threadIdx.x;
	if (idx < size){
		w[idx] = complex_exp(-2*_PI*idx/size);
	}
}

extern "C"
__global__ void inverse(int size, int power, float2* src){
	int idx = blockIdx.x * gridDim.x + threadIdx.x;
	int idy = blockIdx.y * gridDim.y + threadIdx.y;
	if (idx < size && idy < size){
		
		int indexy = idy;
		int indexx = idx;
		int newrow = 0;
		int newcol = 0;
		for (int l = power; l >= 0; l--){
			int k = 1 << l;
			if (k <= indexy){
				indexy -= k;
				int delta = 1 << (power-l);
				newrow += delta;
			}
			if (k <= indexx){
				indexx -= k;
				int delta = 1 << (power-l);
				newcol += delta;
			}
		}
		int rowcol = idy*size + idx;
		int newrownewcol = newrow*size + newcol;
		if (idy > newrow){
			float2 t = src[rowcol];
			src[rowcol] = src[newrownewcol];
			src[newrownewcol] = t;
		} else if (idy == newrow && idx > newcol){
			float2 t = src[rowcol];
			src[rowcol] = src[newrownewcol];
			src[newrownewcol] = t;
		}
	}
}

extern "C"
__global__ void inverseRows(int size, float2* src){
	int blockSize = blockDim.x * blockDim.y;
	int block = blockIdx.y * gridDim.x + blockIdx.x;
	int idx = (blockSize * block) + threadIdx.y * blockDim.x + threadIdx.x;
	int col = idx;
	if (idx < size){
		int icol = 0;
		int jcol = 0;
		int dl = size/2;
		int st = size - 1;
		int j = 0;
		for (int i = 0; i < st; i++){
			if (i < j){
				icol = i*size + col;
				jcol = j*size + col;
				float2 s0 = src[icol];
				src[icol] = src[jcol];
				src[jcol] = s0;
			}
			int k = dl;
			while (k <= j){
				j = j - k;
				k = k >> 1;
			}
			j = j + k;
		}
	}
}

extern "C"
__global__ void inverseCols(int size, float2* src){
	int blockSize = blockDim.x * blockDim.y;
	int block = blockIdx.y * gridDim.x + blockIdx.x;
	int idx = (blockSize * block) + threadIdx.y * blockDim.x + threadIdx.x;
	int row = idx;
	if (idx < size){
		int rowi = 0;
		int rowj = 0;
		int dl = size/2;
		int st = size - 1;
		int j = 0;
		for (int i = 0; i < st; i++){
			if (i < j){
				rowi = row*size + i;
				rowj = row*size + j;
				float2 s0 = src[rowi];
				src[rowi] = src[rowj];
				src[rowj] = s0;
			}
			int k = dl;
			while (k <= j){
				j = j - k;
				k = k >> 1;
			}
			j = j + k;
		}
	}
}
